#ifndef COMMON_MATHS_MATRIX44_H
#define COMMON_MATHS_MATRIX44_H


#include "Vector4.h"


class Matrix44
{
public:
    enum ZeroTag { kZero };
    enum IdentityTag { kIdentity };
    
                        Matrix44();
                        Matrix44(const Matrix44& mat);
                        Matrix44(Vector4Arg row0, Vector4Arg row1, Vector4Arg row2, Vector4Arg row3);
                        Matrix44(const ZeroTag&);
                        Matrix44(const IdentityTag&);
                        Matrix44(Vector4Arg axis, float angle);

    static Matrix44     Scale(float scale);
    static Matrix44     Scale(Vector3Arg scale);
    static Matrix44     Scale(Vector4Arg scale);
    static Matrix44     RotateX(float angle);
    static Matrix44     RotateY(float angle);
    static Matrix44     RotateZ(float angle);
    static Matrix44     Translate(Vector3Arg translation);

    float               operator[](size_t index) const;
    float&              operator[](size_t index);
    
    const float*        GetElements() const;
    float*              GetElements();

    Vector4Arg          GetRow(size_t rowIdx) const;
    Vector4             GetCol(size_t colIdx) const;
    void                SetRow(size_t rowIdx, Vector4Arg v);
    void                SetCol(size_t colIdx, Vector4Arg v);
    Matrix44&           SetTranslation(Vector3Arg v);

    Vector3Arg          GetXAxis() const;
    Vector3Arg          GetYAxis() const;
    Vector3Arg          GetZAxis() const;
    Vector3Arg          GetTranslation() const;

    Matrix44&           operator=(const Matrix44& mat);
    Matrix44&           operator=(const ZeroTag&);
    Matrix44&           operator=(const IdentityTag&);

    Matrix44            operator+(const Matrix44& mat) const;
    Matrix44            operator-(const Matrix44& mat) const;
    Matrix44&           operator+=(const Matrix44& mat);
    Matrix44&           operator-=(const Matrix44& mat);
    Matrix44            operator-() const;

    Matrix44            operator*(float v) const;
    Matrix44&           operator*=(float v);
    Matrix44            operator*(const Matrix44& mat) const;
    Matrix44&           operator*=(const Matrix44& mat);

    friend Matrix44     MulTransformMatrix(const Matrix44& trans, const Matrix44& mat);
    friend Vector3      MulPointTransform(Vector3Arg point, const Matrix44& mat);
    friend Vector3      MulVectorTransform(Vector3Arg vec, const Matrix44& mat);
    friend Vector4      MulVectorMatrix(Vector4Arg vec, const Matrix44& mat);

    friend Matrix44     Abs(const Matrix44& mat);
    friend Matrix44     Transpose(const Matrix44& mat);
    friend Matrix44     TransformInverseRT(const Matrix44& mat);
    friend Matrix44     TransformInverseRTS(const Matrix44& mat);
    friend Matrix44     TransformInverseTransposeRT(const Matrix44& mat);
    friend Matrix44     TransformInverseTransposeRTS(const Matrix44& mat);
    friend Matrix44     operator*(float f, const Matrix44& mat);

private:
    Vector4             row[4];
};


inline Matrix44::Matrix44()
{
}

inline Matrix44::Matrix44(const Matrix44& mat)
{
    row[0] = mat.row[0];
    row[1] = mat.row[1];
    row[2] = mat.row[2];
    row[3] = mat.row[3];
}

inline Matrix44::Matrix44(Vector4Arg row0, Vector4Arg row1, Vector4Arg row2, Vector4Arg row3)
{
    row[0] = row0;
    row[1] = row1;
    row[2] = row2;
    row[3] = row3;
}

inline Matrix44::Matrix44(const ZeroTag&)
{
    const Vector4 zero = Vector4(0.0f, 0.0f, 0.0f, 0.0f);
    row[0] = zero;
    row[1] = zero;
    row[2] = zero;
    row[3] = zero;
}

inline Matrix44::Matrix44(const IdentityTag&)
{
    row[0] = Vector4(1.0f, 0.0f, 0.0f, 0.0f);
    row[1] = Vector4(0.0f, 1.0f, 0.0f, 0.0f);
    row[2] = Vector4(0.0f, 0.0f, 1.0f, 0.0f);
    row[3] = Vector4(0.0f, 0.0f, 0.0f, 1.0f);
}

inline Matrix44::Matrix44(Vector4Arg axis, float angle)
{
    float x = axis.x;
    float y = axis.y;
    float z = axis.z;
    float c = cosf(angle);
    float s = sinf(angle);

    row[0] = Vector4(x*x*(1-c) + c, y*x*(1-c) + z*s, x*z*(1-c) - y*s, 0.0f);
    row[1] = Vector4(x*y*(1-c) - z*s, y*y*(1-c) + c, y*z*(1-c) + x*s, 0.0f);
    row[2] = Vector4(x*z*(1-c) + y*s, y*z*(1-c) - x*s, z*z*(1-c) + c, 0.0f);
    row[3] = Vector4(0.0f, 0.0f, 0.0f, 1.0f);
}

inline float Matrix44::operator[](size_t index) const
{
    const float* p = &row[0].x;
    return p[index];
}

inline float& Matrix44::operator[](size_t index)
{
    float* p = &row[0].x;
    return p[index];
}

inline const float* Matrix44::GetElements() const
{
    return (const float*)&row[0].x;
}

inline float* Matrix44::GetElements()
{
    return (float*)&row[0].x;
}

inline Vector4Arg Matrix44::GetRow(size_t rowIdx) const
{
    return row[rowIdx];
}

inline Vector4 Matrix44::GetCol(size_t colIdx) const
{
    const float* p = &row[0].x + colIdx;
    return Vector4(p[0], p[4], p[8], p[12]);
}

inline Vector3Arg Matrix44::GetXAxis() const
{
    return *(Vector3*)&row[0];
}

inline Vector3Arg Matrix44::GetYAxis() const
{
    return *(Vector3*)&row[1];
}

inline Vector3Arg Matrix44::GetZAxis() const
{
    return *(Vector3*)&row[2];
}

inline Vector3Arg Matrix44::GetTranslation() const
{
    return *(Vector3*)&row[3];
}

inline void Matrix44::SetRow(size_t rowIdx, Vector4Arg v)
{
    row[rowIdx] = v;
}

inline Matrix44& Matrix44::SetTranslation(Vector3Arg v)
{
    row[3] = Vector4(v, 1.0f);
    return *this;
}

inline void Matrix44::SetCol(size_t colIdx, Vector4Arg v)
{
    float* p = &row[0].x + colIdx;
    p[0] = v.x;
    p[4] = v.y;
    p[8] = v.z;
    p[12] = v.w;
}

inline Matrix44 Matrix44::Scale(float scale)
{
    return Scale(Vector4(scale, scale, scale, 1.0f));
}

inline Matrix44 Matrix44::Scale(Vector3Arg scale)
{
    return Scale(Vector4(scale, 1.0f));
}

inline Matrix44 Matrix44::Scale(Vector4Arg scale)
{
    return Matrix44(
        Vector4(scale.x, 0.0f, 0.0f, 0.0f),
        Vector4(0.0f, scale.y, 0.0f, 0.0f),
        Vector4(0.0f, 0.0f, scale.z, 0.0f),
        Vector4(0.0f, 0.0f, 0.0f, scale.w));
}

inline Matrix44 Matrix44::RotateX(float angle)
{
    float c = cosf(angle);
    float s = sinf(angle);
    return Matrix44(Vector4(1.0f, 0.0f, 0.0f, 0.0f),
        Vector4(0.0f,    c,    s, 0.0f),
        Vector4(0.0f,   -s,    c, 0.0f),
        Vector4(0.0f, 0.0f, 0.0f, 1.0f));
}

inline Matrix44 Matrix44::RotateY(float angle)
{
    float c = cosf(angle);
    float s = sinf(angle);
    return Matrix44(
        Vector4(  c, 0.0f,   -s, 0.0f),
        Vector4(0.0f, 1.0f, 0.0f, 0.0f),
        Vector4(  s, 0.0f,    c, 0.0f),
        Vector4(0.0f, 0.0f, 0.0f, 1.0f));
}

inline Matrix44 Matrix44::RotateZ(float angle)
{
    float c = cosf(angle);
    float s = sinf(angle);
    return Matrix44(
        Vector4(  c,    s, 0.0f, 0.0f),
        Vector4( -s,    c, 0.0f, 0.0f),
        Vector4(0.0f, 0.0f, 1.0f, 0.0f),
        Vector4(0.0f, 0.0f, 0.0f, 1.0f));
}

inline Matrix44 Matrix44::Translate(Vector3Arg translation)
{
    return Matrix44(
        Vector4(1, 0, 0, 0),
        Vector4(0, 1, 0, 0),
        Vector4(0, 0, 1, 0),
        Vector4(translation, 1));
}

inline Matrix44& Matrix44::operator=(const Matrix44& mat)
{
    row[0] = mat.row[0];
    row[1] = mat.row[1];
    row[2] = mat.row[2];
    row[3] = mat.row[3];
    return *this;
}

inline Matrix44& Matrix44::operator=(const ZeroTag&)
{
    const Vector4 zero = Vector4(0.0f, 0.0f, 0.0f, 0.0f);
    row[0] = zero;
    row[1] = zero;
    row[2] = zero;
    row[3] = zero;
    return *this;
}

inline Matrix44& Matrix44::operator=(const IdentityTag&)
{
    row[0] = Vector4(1.0f, 0.0f, 0.0f, 0.0f);
    row[1] = Vector4(0.0f, 1.0f, 0.0f, 0.0f);
    row[2] = Vector4(0.0f, 0.0f, 1.0f, 0.0f);
    row[3] = Vector4(0.0f, 0.0f, 0.0f, 1.0f);
    return *this;
}

inline Matrix44 Matrix44::operator+(const Matrix44& mat) const
{
    return Matrix44(
        row[0] + mat.row[0],
        row[1] + mat.row[1],
        row[2] + mat.row[2],
        row[3] + mat.row[3]
    );
}

inline Matrix44 Matrix44::operator-(const Matrix44& mat) const
{
    return Matrix44(
        row[0] - mat.row[0],
        row[1] - mat.row[1],
        row[2] - mat.row[2],
        row[3] - mat.row[3]
    );
}

inline Matrix44& Matrix44::operator+=(const Matrix44& mat)
{
    row[0] += mat.row[0];
    row[1] += mat.row[1];
    row[2] += mat.row[2];
    row[3] += mat.row[3];
    return *this;
}

inline Matrix44& Matrix44::operator-=(const Matrix44& mat)
{
    row[0] -= mat.row[0];
    row[1] -= mat.row[1];
    row[2] -= mat.row[2];
    row[3] -= mat.row[3];
    return *this;
}

inline Matrix44 Matrix44::operator-() const
{
    return Matrix44(-row[0], -row[1], -row[2], -row[3]);
}

inline Matrix44 Matrix44::operator*(float v) const
{
    return Matrix44(
        row[0] * v,
        row[1] * v,
        row[2] * v,
        row[3] * v
    );
}

inline Matrix44& Matrix44::operator*=(float v)
{
    row[0] *= v;
    row[1] *= v;
    row[2] *= v;
    row[3] *= v;
    return *this;
}

inline Matrix44 Matrix44::operator*(const Matrix44& mat) const
{
    return Matrix44(
        MulVectorMatrix(row[0], mat),
        MulVectorMatrix(row[1], mat),
        MulVectorMatrix(row[2], mat),
        MulVectorMatrix(row[3], mat)
    );
}

inline Matrix44& Matrix44::operator*=(const Matrix44& mat)
{
    *this = (*this) * mat;
    return *this;
}

inline Matrix44 MulTransformMatrix(const Matrix44& trans, const Matrix44& mat)
{
    const Vector4 a =
        mat.row[0] * trans.row[0].x +
        mat.row[1] * trans.row[0].y +
        mat.row[2] * trans.row[0].z;

    const Vector4 b =
        mat.row[0] * trans.row[1].x +
        mat.row[1] * trans.row[1].y +
        mat.row[2] * trans.row[1].z;
    
    const Vector4 c =
        mat.row[0] * trans.row[2].x +
        mat.row[1] * trans.row[2].y +
        mat.row[2] * trans.row[2].z;

    const Vector4 d =
        mat.row[0] * trans.row[3].x +
        mat.row[1] * trans.row[3].y +
        mat.row[2] * trans.row[3].z +
        mat.row[3];
    
    return Matrix44(a, b, c, d);
}

inline Vector3 MulPointTransform(Vector3Arg point, const Matrix44& mat)
{
    const Vector3 a = mat.GetXAxis() * point.x;
    const Vector3 b = mat.GetYAxis() * point.y;
    const Vector3 c = mat.GetZAxis() * point.z;
    const Vector3 d = mat.GetTranslation();
    return (a+b) + (c+d);
}

inline Vector3 MulVectorTransform(Vector3Arg vec, const Matrix44& mat)
{
    const Vector3 a = mat.GetXAxis() * vec.x;
    const Vector3 b = mat.GetYAxis() * vec.y;
    const Vector3 c = mat.GetZAxis() * vec.z;
    return a + b + c;
}

inline Vector4 MulVectorMatrix(Vector4Arg vec, const Matrix44& mat)
{
    const Vector4 a = mat.row[0] * vec.x;
    const Vector4 b = mat.row[1] * vec.y;
    const Vector4 c = mat.row[2] * vec.z;
    const Vector4 d = mat.row[3] * vec.w;
    return (a+b) + (c+d);
}

inline Matrix44 Abs(const Matrix44& mat)
{
    return Matrix44(Abs(mat.row[0]), Abs(mat.row[1]), Abs(mat.row[2]), Abs(mat.row[3]));
}

inline Matrix44 Transpose(const Matrix44& mat)
{
    Matrix44 result;
    result.row[0] = Vector4(mat.row[0].x, mat.row[1].x, mat.row[2].x, mat.row[3].x);
    result.row[1] = Vector4(mat.row[0].y, mat.row[1].y, mat.row[2].y, mat.row[3].y);
    result.row[2] = Vector4(mat.row[0].z, mat.row[1].z, mat.row[2].z, mat.row[3].z);
    result.row[3] = Vector4(mat.row[0].w, mat.row[1].w, mat.row[2].w, mat.row[3].w);
    return result;
}

// This assumes that the input matrix represents an affine transform. The function
// behaves as though the 4th column is (0,0,0,1), and will return a fully initialised Matrix44
// This should be used when 3x3 rotation component is orthogonal, as it's faster.
// Dosn't handle matrices with scaling
inline Matrix44 TransformInverseRT(const Matrix44& mat)
{
    Matrix44 result;
    result.row[0] = Vector4(mat.row[0].x, mat.row[1].x, mat.row[2].x, 0.0f);
    result.row[1] = Vector4(mat.row[0].y, mat.row[1].y, mat.row[2].y, 0.0f);
    result.row[2] = Vector4(mat.row[0].z, mat.row[1].z, mat.row[2].z, 0.0f);
    
    const Vector4 t = mat.row[3];
    result.row[3].x = -(result.row[0].x*t.x + result.row[1].x*t.y + result.row[2].x*t.z);
    result.row[3].y = -(result.row[0].y*t.x + result.row[1].y*t.y + result.row[2].y*t.z);
    result.row[3].z = -(result.row[0].z*t.x + result.row[1].z*t.y + result.row[2].z*t.z);
    result.row[3].w = 1.0f;
    
    return result;
}

// This assumes that the input matrix represents an affine transform. The function
// behaves as though the 4th column is (0,0,0,1), and will return a fully initialised Matrix44
// This should be used when 3x3 rotation component is orthogonal, as it's faster.
// Handles matrices with scaling
inline Matrix44 TransformInverseRTS(const Matrix44& mat)
{
    float s = 1.0f / (mat[0]*mat[0] + mat[4]*mat[4] + mat[8]*mat[8]);
    Matrix44 result;
    result.row[0] = s * Vector4(mat.row[0].x, mat.row[1].x, mat.row[2].x, 0.0f);
    result.row[1] = s * Vector4(mat.row[0].y, mat.row[1].y, mat.row[2].y, 0.0f);
    result.row[2] = s * Vector4(mat.row[0].z, mat.row[1].z, mat.row[2].z, 0.0f);
    
    const Vector4 t = mat.row[3];
    result.row[3].x = -(result.row[0].x*t.x + result.row[1].x*t.y + result.row[2].x*t.z);
    result.row[3].y = -(result.row[0].y*t.x + result.row[1].y*t.y + result.row[2].y*t.z);
    result.row[3].z = -(result.row[0].z*t.x + result.row[1].z*t.y + result.row[2].z*t.z);
    result.row[3].w = 1.0f;
    
    return result;
}

// This assumes that the input matrix represents an affine transform. The function
// behaves as though the 4th column is (0,0,0,1), and will return a fully initialised Matrix44
// This should be used when 3x3 rotation component is orthogonal, as it's faster.
// Dosn't handle matrices with scaling
inline Matrix44 TransformInverseTransposeRT(const Matrix44& mat)
{
    Matrix44 result;
    float a = -Dot3(mat.row[0], mat.row[3]);
    float b = -Dot3(mat.row[1], mat.row[3]);
    float c = -Dot3(mat.row[2], mat.row[3]);
    result.row[0] = Vector4(mat.row[0].x, mat.row[0].y, mat.row[0].z, a);
    result.row[1] = Vector4(mat.row[1].x, mat.row[1].y, mat.row[1].z, b);
    result.row[2] = Vector4(mat.row[2].x, mat.row[2].y, mat.row[2].z, c);
    result.row[3] = Vector4(0.0f, 0.0f, 0.0f, 1.0f);
    return result;
}

// This assumes that the input matrix represents an affine transform. The function
// behaves as though the 4th column is (0,0,0,1), and will return a fully initialised Matrix44
// This should be used when 3x3 rotation component is orthogonal, as it's faster.
// Handles matrices with scaling
inline Matrix44 TransformInverseTransposeRTS(const Matrix44& mat)
{
    float s = 1.0f / (mat[0]*mat[0] + mat[4]*mat[4] + mat[8]*mat[8]);
    Matrix44 result;
    float a = -Dot3(mat.row[0], mat.row[3]);
    float b = -Dot3(mat.row[1], mat.row[3]);
    float c = -Dot3(mat.row[2], mat.row[3]);
    result.row[0] = s * Vector4(mat.row[0].x, mat.row[0].y, mat.row[0].z, a);
    result.row[1] = s * Vector4(mat.row[1].x, mat.row[1].y, mat.row[1].z, b);
    result.row[2] = s * Vector4(mat.row[2].x, mat.row[2].y, mat.row[2].z, c);
    result.row[3] = Vector4(0.0f, 0.0f, 0.0f, 1.0f);
    return result;
}


inline Matrix44 operator*(float f, const Matrix44& mat)
{
    return mat * f;
}


#endif
